# -*- codeing = utf-8 -*-
# @Time : 2022/2/18 13:31
# @File : model.py
# @Software : PyCharm
from torch import nn
import torch.nn.functional as F

import database

# 卷积神经网络模型
class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        # 第一层卷积->激活->池化->Dropout
        self.conv1 = nn.Sequential(
            # 卷积计算
            # in_channels: 输入的RGB图像为3层
            # out_channels： 输出的16层,
            # kernel_size: 卷积核为5*5
            # stride步长为1*1  每次卷积计算自后像右（下）移动一位
            # padding: 图像周围加2*2的边框，卷积核为5*5，图片加上2*2的边框之后卷积计算得到结果还是32*32大小的图片（32-5+2+2)
            nn.Conv2d(
                in_channels=3,
                out_channels=16,
                kernel_size=5,
                stride=1,
                padding=2,
            ),
            nn.ReLU(),  # 激活函数
            nn.MaxPool2d(kernel_size=2), # 池化层2*2大小，32*32大小的图片经过2*2的池化后就是 16*16（32/2,32/2)大小的图片
            nn.Dropout(0.2) # 随机丢弃0.2的神经元权重，dropout是深度神经网络防止过度拟合很好的方式
        )
        # 第二层卷积->激活->池化->Dropout
        self.conv2 = nn.Sequential(
            # 同上
            # 第一层出来16，第二层16进，32出，还原图像
            nn.Conv2d(16, 32, 5, 1, 2),
            nn.ReLU(),
            nn.MaxPool2d(2),
            nn.Dropout(0.2)
        )
        # 全连接层
        # 32*32的图片经过两次池化之后大小为8*8,上一层卷积输出是32层，所以总共的大小是32*8*8个神经源，
        # 输出分类由要识别的人脸数量决定
        self.out = nn.Linear(32 * 8 * 8, 3)

    # pytorch神经网络推理过程，将前面的计算串联在一起计算结果
    def forward(self, x):
        x = self.conv1(x)
        x = self.conv2(x)
        x = x.view(x.size(0), -1)
        x = self.out(x)
        # 对结果进行log+softmax输出
        return F.log_softmax(x, dim=1)